
"""
数据预处理
"""
import multiprocessing
import mindspore.common.dtype as mstype
import mindspore.dataset as ds
import mindspore.dataset.vision.c_transforms as C
import mindspore.dataset.transforms.c_transforms as C2
from mindspore.communication.management import init, get_rank, get_group_size

"""
cifar10 dataset for resnet50
"""
def create_dataset(dataset_path ,do_train ,repeat_num = 1 ,batch_size= 32 ,train_image_size = 224,
                   distribute=False ,enable_cache=False ,cache_session_id=None):
    device_num, rank_id = _get_rank_info(distribute)
    ds.config.set_prefetch_size(64)
    if device_num == 1:
        data_set = ds.Cifar10Dataset(dataset_path, num_parallel_workers=get_num_parallel_workers(12), shuffle=True)
    else:
        data_set = ds.Cifar10Dataset(dataset_path, num_parallel_workers=get_num_parallel_workers(12), shuffle=True,
                                     num_shards=device_num, shard_id=rank_id)

    # define map operations
    trans = []
    if do_train:
        trans += [
            C.RandomCrop((32, 32), (4, 4, 4, 4)),
            C.RandomHorizontalFlip(prob=0.5)
        ]

    trans += [
        C.Resize((train_image_size, train_image_size)),
        C.Rescale(1.0 / 255.0, 0.0),
        C.Normalize([0.4914, 0.4822, 0.4465], [0.2023, 0.1994, 0.2010]),
        C.HWC2CHW()
    ]

    type_cast_op = C2.TypeCast(mstype.int32)
    data_set = data_set.map(operations=type_cast_op, input_columns="label",
                            num_parallel_workers=get_num_parallel_workers(8))
    data_set = data_set.map(operations=trans, input_columns="image",
                                num_parallel_workers=get_num_parallel_workers(8))

    # apply batch operations
    data_set = data_set.batch(batch_size, drop_remainder=True)
    # apply dataset repeat operation
    data_set = data_set.repeat(repeat_num)

    return data_set

def _get_rank_info(distribute):
    if distribute:
        init()
        rank_id = get_rank()
        device_num = get_group_size()
    else:
        rank_id = 0
        device_num = 1
    return device_num ,rank_id

def get_num_parallel_workers(num_parallel_workers):

    cores = multiprocessing.cpu_count()
    if isinstance(num_parallel_workers ,int):
        if cores < num_parallel_workers:
            num_parallel_workers = cores

    else:
        num_parallel_workers = min(cores ,8)
    return num_parallel_workers
